# src/evaluate.py

import torch
from torch.utils.data import DataLoader
from data_loader import ECGDataset
from model_wavelet import ECG_Autoencoder_Classifier_Lstm, ECG_Autoencoder_Classifier_Transformer, ECG_Autoencoder_Classifier_Mamba, ECG_Autoencoder_Classifier_Cnn# 导入新的模型类
from model_time import ChannelWiseCNN
from sklearn.metrics import classification_report, roc_auc_score
import pandas as pd
import os
from datetime import datetime
import numpy as np
import json  # 添加导入

def evaluate_model(model, dataloader, device):
    # 创建时间戳文件夹
    timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
    result_dir = os.path.join('results/evaluation', timestamp)
    os.makedirs(result_dir, exist_ok=True)

    model.eval()
    preds_list = []
    labels_list = []
    sample_ids = []  # 新增：用于存储样本编号

    with torch.no_grad():
        for batch in dataloader:  # 修改这里以适应新的返回值
            inputs, labels, ids = batch  # 新增：接收样本编号
            inputs = inputs.to(device)
            labels = labels.to(device).float()

            logits = model(inputs)
            # 应用 Sigmoid 获取概率
            probs = torch.sigmoid(logits)
            preds = (probs >= 0.55).float()

            preds_list.extend(preds.cpu().numpy())
            labels_list.extend(labels.cpu().numpy())
            sample_ids.extend(ids)  # 新增：收集样本编号

    if len(labels_list) == 0 or len(preds_list) == 0:
        print("Error: No samples found in the test dataset or no predictions were made.")
        return

    # 将标签和预测结果展平
    preds_list = np.array(preds_list).flatten()
    labels_list = np.array(labels_list).flatten()

    # 生成评估报告
    report = classification_report(labels_list, preds_list, output_dict=True)
    report_text = classification_report(labels_list, preds_list)
    try:
        auc = roc_auc_score(labels_list, preds_list)
    except ValueError:
        auc = 0.0  # 如果无法计算 AUC（例如只有一个类别），则设为 0.0
        print('ROC-AUC Score cannot be calculated (only one class present in y_true).')

    # 打印结果
    print(report_text)
    print(f'ROC-AUC Score: {auc:.4f}')

    # 保存评估报告到 txt 文件
    with open(os.path.join(result_dir, 'evaluation_report.txt'), 'w') as f:
        f.write(report_text)
        f.write(f'\nROC-AUC Score: {auc:.4f}')

    # 保存评估结果到 CSV 文件
    report_df = pd.DataFrame(report).transpose()
    report_df['roc_auc'] = auc
    report_df.to_csv(os.path.join(result_dir, 'evaluation_results.csv'), index=True)

    # 保存预测结果到 CSV 文件
    correct_list = [int(pred == label) for pred, label in zip(preds_list, labels_list)]
    results_df = pd.DataFrame({
        'Sample_ID': sample_ids,  # 新增：添加样本编号列
        'True_Labels': labels_list,
        'Predictions': preds_list,
        'Correct': correct_list
    })
    
    # 添加排序逻辑：将sample_id转换为数字进行排序
    results_df['sort_key'] = results_df['Sample_ID'].str.extract('(\d+)').astype(int)
    results_df = results_df.sort_values('sort_key')
    results_df = results_df.drop('sort_key', axis=1)  # 删除辅助排序列
    
    results_df.to_csv(os.path.join(result_dir, 'predictions.csv'), index=False)
    # 保存模型权重
    torch.save(model.state_dict(), os.path.join(result_dir, 'model_weights.pth'))
    
    # 保存模型设置
    model_config = {
        'input_dim': input_dim,       # 修改为正确的属性路径
        'hidden_dim': hidden_dim,
        'latent_dim': latent_dim,
        'num_layers': num_layers
    }
    with open(os.path.join(result_dir, 'model_config.json'), 'w') as f:
        json.dump(model_config, f)

if __name__ == '__main__':
    test_data_dir = 'data/processed/test_data'
    batch_size = 32
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    # 数据加载
    test_dataset = ECGDataset(test_data_dir)
    test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)

    # 模型加载
    input_dim = 1  # 与数据加载器中 ecg 的特征维度一致
    hidden_dim = 256  # 确保与 train.py 中一致
    latent_dim = 128   # 确保与 train.py 中一致
    num_layers = 8   # 确保与 train.py 中一致

    model = ECG_Autoencoder_Classifier_Cnn(input_dim, hidden_dim, latent_dim, num_layers)
    # 加载模型参数，并指定 map_location，以确保在 CPU 或 GPU 上正确加载
    #model.load_state_dict(torch.load('results/evaluation/20241215_175027/model_weights.pth', map_location=device)['model_state_dict'])
    model.load_state_dict(torch.load('results/final_model/Cnn/model_weights.pth', map_location=device))
    model = model.to(device)

    # 模型评估
    evaluate_model(model, test_loader, device)

